# Created by Gao Song on 2019-04-17

import numpy as np


def contract_range(sym_near, sym_far):
    sym_near, exchange = sym_near.split('.')
    sym_far, exchange = sym_far.split('.')

    if exchange == 'CZCE':
        y1 = int(sym_near[-3:-2])
        y2 = int(sym_far[-3:-2])
        if y2 == 0:
            y2 = 10
    else:
        y1 = int(sym_near[-4:-2])
        y2 = int(sym_far[-4:-2])

    m1 = int(sym_near[-2:])
    m2 = int(sym_far[-2:])
    return 12 * (y2 - y1) + (m2 - m1)

def reg_basis(log_price, reg_w):
    syms = sorted(reg_w[reg_w > 0].index.to_list())

    if len(syms) < 2:
        return np.nan
    x = [contract_range(syms[0], sym) for sym in syms]
    y = log_price[syms].values
    w = reg_w[syms].values
    p = np.poly1d(np.polyfit(x, y, 1, w=w))

    return p[1]
